from .flexflow_cffi_header import ffc, ffi
from enum import Enum

class DataType(Enum):
  BOOL = ffc.FLEXFLOW_DATATYPE_BOOL
  INT32 = ffc.FLEXFLOW_DATATYPE_INT32
  INT64 = ffc.FLEXFLOW_DATATYPE_INT64
  HALF = ffc.FLEXFLOW_DATATYPE_HALF
  FLOAT = ffc.FLEXFLOW_DATATYPE_FLOAT
  DOUBLE = ffc.FLEXFLOW_DATATYPE_DOUBLE

class Activation(Enum):
  RELU = ffc.FLEXFLOW_ACTIVATION_RELU
  SIGMOID = ffc.FLEXFLOW_ACTIVATION_SIGMOID
  TANH = ffc.FLEXFLOW_ACTIVATION_TANH
  GELU = ffc.FLEXFLOW_ACTIVATION_GELU
  NONE = ffc.FLEXFLOW_ACTIVATION_NONE

class PoolOp(Enum):
  MAX = ffc.FLEXFLOW_POOL_OP_MAX
  AVG = ffc.FLEXFLOW_POOL_OP_AVG

class ParamSyncType(Enum):
  PARAM_SERVER = ffc.FLEXFLOW_PARAM_SYNC_PARAMETER_SERVER
  NCCL = ffc.FLEXFLOW_PARAM_SYNC_NCCL

class AggregateOp(Enum):
  SUM = ffc.FLEXFLOW_AGGREGATE_OP_SUM
  AVG = ffc.FLEXFLOW_AGGREGATE_OP_AVG

class OpType(Enum):
  NOOP = ffc.FLEXFLOW_OP_TYPE_NOOP
  INPUT = ffc.FLEXFLOW_OP_TYPE_INPUT
  WEIGHT = ffc.FLEXFLOW_OP_TYPE_WEIGHT
  CONV2D = ffc.FLEXFLOW_OP_TYPE_CONV2D
  DROPOUT = ffc.FLEXFLOW_OP_TYPE_DROPOUT
  LINEAR = ffc.FLEXFLOW_OP_TYPE_LINEAR
  BATCHMATMUL = ffc.FLEXFLOW_OP_TYPE_BATCHMATMUL
  POOL2D = ffc.FLEXFLOW_OP_TYPE_POOL2D
  SCALAR_MULTIPLY = ffc.FLEXFLOW_OP_TYPE_SCALAR_MULTIPLY
  SCALAR_ADD = ffc.FLEXFLOW_OP_TYPE_SCALAR_ADD
  SCALAR_FLOOR_DIV = ffc.FLEXFLOW_OP_TYPE_SCALAR_FLOOR_DIV
  SCALAR_TRUE_DIV = ffc.FLEXFLOW_OP_TYPE_SCALAR_TRUE_DIV
  SCALAR_SUB = ffc.FLEXFLOW_OP_TYPE_SCALAR_SUB
  RELU = ffc.FLEXFLOW_OP_TYPE_RELU
  IDENTITY = ffc.FLEXFLOW_OP_TYPE_IDENTITY
  SIGMOID = ffc.FLEXFLOW_OP_TYPE_SIGMOID
  TANH = ffc.FLEXFLOW_OP_TYPE_TANH
  ELU = ffc.FLEXFLOW_OP_TYPE_ELU
  FLAT = ffc.FLEXFLOW_OP_TYPE_FLAT
  SOFTMAX = ffc.FLEXFLOW_OP_TYPE_SOFTMAX
  BATCHNORM = ffc.FLEXFLOW_OP_TYPE_BATCHNORM
  CONCAT = ffc.FLEXFLOW_OP_TYPE_CONCAT
  SPLIT = ffc.FLEXFLOW_OP_TYPE_SPLIT
  EMBEDDING = ffc.FLEXFLOW_OP_TYPE_EMBEDDING
  CACHE = ffc.FLEXFLOW_OP_TYPE_CACHE
  RESHAPE = ffc.FLEXFLOW_OP_TYPE_RESHAPE
  REVERSE = ffc.FLEXFLOW_OP_TYPE_REVERSE
  TRANSPOSE = ffc.FLEXFLOW_OP_TYPE_TRANSPOSE
  EW_ADD = ffc.FLEXFLOW_OP_TYPE_EW_ADD
  EW_MUL = ffc.FLEXFLOW_OP_TYPE_EW_MUL
  MATMUL = ffc.FLEXFLOW_OP_TYPE_MATMUL
  MUL = ffc.FLEXFLOW_OP_TYPE_MUL
  ENLARGE = ffc.FLEXFLOW_OP_TYPE_ENLARGE
  SQUEEZE = ffc.FLEXFLOW_OP_TYPE_SQUEEZE 
  UNSQUEEZE = ffc.FLEXFLOW_OP_TYPE_UNSQUEEZE 
  EW_SUB = ffc.FLEXFLOW_OP_TYPE_EW_SUB   
  EW_DIV = ffc.FLEXFLOW_OP_TYPE_EW_DIV   
  EW_EQUAL = ffc.FLEXFLOW_OP_TYPE_EW_EQUAL 
  EW_GREATER = ffc.FLEXFLOW_OP_TYPE_EW_GREATER 
  EW_LESS = ffc.FLEXFLOW_OP_TYPE_EW_LESS 
  EW_MAX = ffc.FLEXFLOW_OP_TYPE_EW_MAX  
  EW_MIN = ffc.FLEXFLOW_OP_TYPE_EW_MIN  
  REDUCE_ARGMAX = ffc.FLEXFLOW_OP_TYPE_REDUCE_ARGMAX 
  REDUCE_ARGMIN = ffc.FLEXFLOW_OP_TYPE_REDUCE_ARGMIN 
  REDUCE_MAX = ffc.FLEXFLOW_OP_TYPE_REDUCE_MAX 
  REDUCE_MEAN = ffc.FLEXFLOW_OP_TYPE_REDUCE_MEAN 
  REDUCE_MIN = ffc.FLEXFLOW_OP_TYPE_REDUCE_MIN 
  REDUCE_PROD = ffc.FLEXFLOW_OP_TYPE_REDUCE_PROD 
  REDUCE_SUM = ffc.FLEXFLOW_OP_TYPE_REDUCE_SUM 
  PAD = ffc.FLEXFLOW_OP_TYPE_PAD 
  SHAPE = ffc.FLEXFLOW_OP_TYPE_SHAPE 
  SIZE = ffc.FLEXFLOW_OP_TYPE_SIZE  
  TOPK = ffc.FLEXFLOW_OP_TYPE_TOPK  
  WHERE = ffc.FLEXFLOW_OP_TYPE_WHERE 
  CEIL = ffc.FLEXFLOW_OP_TYPE_CEIL  
  CAST = ffc.FLEXFLOW_OP_TYPE_CAST  
  EXP = ffc.FLEXFLOW_OP_TYPE_EXP   
  ROUND = ffc.FLEXFLOW_OP_TYPE_ROUND 
  LOG = ffc.FLEXFLOW_OP_TYPE_LOG   
  LOGICAL_NOT = ffc.FLEXFLOW_OP_TYPE_LOGICAL_NOT 
  SQRT = ffc.FLEXFLOW_OP_TYPE_SQRT 
  SIN = ffc.FLEXFLOW_OP_TYPE_SIN  
  COS = ffc.FLEXFLOW_OP_TYPE_COS  
  LEAKYRELU = ffc.FLEXFLOW_OP_TYPE_LEAKYRELU
  SLICE = ffc.FLEXFLOW_OP_TYPE_SLICE  
  RESIZE = ffc.FLEXFLOW_OP_TYPE_RESIZE 
  PRELU = ffc.FLEXFLOW_OP_TYPE_PRELU  
  GELU = ffc.FLEXFLOW_OP_TYPE_GELU
  MULTIHEAD_ATTENTION = ffc.FLEXFLOW_OP_TYPE_MULTIHEAD_ATTENTION
  FUSED = ffc.FLEXFLOW_OP_TYPE_FUSED 
  RSQRT = ffc.FLEXFLOW_OP_TYPE_RSQRT 
  POW = ffc.FLEXFLOW_OP_TYPE_POW   
  MEAN = ffc.FLEXFLOW_OP_TYPE_MEAN  
  LAYERNORM = ffc.FLEXFLOW_OP_TYPE_LAYERNORM
  GATHER = ffc.FLEXFLOW_OP_TYPE_GATHER 
  BROADCAST = ffc.FLEXFLOW_OP_TYPE_BROADCAST
  REPARTITION = ffc.FLEXFLOW_OP_TYPE_REPARTITION
  COMBINE = ffc.FLEXFLOW_OP_TYPE_COMBINE
  REPLICATE = ffc.FLEXFLOW_OP_TYPE_REPLICATE
  REDUCTION = ffc.FLEXFLOW_OP_TYPE_REDUCTION
  BATCH = ffc.FLEXFLOW_OP_TYPE_BATCH
  PIPELINE = ffc.FLEXFLOW_OP_TYPE_PIPELINE
  FUSED_PARALLEL = ffc.FLEXFLOW_OP_TYPE_FUSED_PARALLEL
